from operator import add
from typing import TypedDict, Annotated

from langchain_core.messages import AnyMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.config import get_stream_writer
from langgraph.constants import START, END
from langgraph.graph import StateGraph

nodes = ["Supervisor", "Imperativequery", "ExploratoryAnalysis", "Dataextraction", "other"]

from zai import ZhipuAiClient
client = ZhipuAiClient(api_key="db2e03cca9654d4dba9b91b123548374.TftHaYZq4ydQPDjP")

class State(TypedDict):
    messages: Annotated[list[AnyMessage], add]
    type: str

def other_node(state: State):
    print("other node")
    writer = get_stream_writer()
    writer({"node", "other_node"})

    return {"messages": [HumanMessage(content="我暂时无法回答这个问题")], "type": "other"}
def supervisor_node(state: State):
    print("Supervisor node")
    writer = get_stream_writer()
    writer({"node", "Supervisor_node"})
    #根据用户的问题，对问题进行分类。分类结果保存到type当中
    prompt = """你是一个专业的SQL分类助手，负责对用户的问题关键词进行分类，并将任务分给其他的Agent执行
                如果用户的问题是和指令式查询有关的，请将任务交给Imperativequery_agent处理，指令式查询:包括事实查询(获取具体数据)、状态查询(当前状态)、简单统计(计数、求和等)。
                如果用户问题是和探索式分析有关的，请将任务交给Exploratoryanalysis_agent处理，探索式分析:包括数据探索(浏览数据)、模式发现(趋势分析)、假设验证(测试假设)。
                如果用户问题是和数据提取有关的，请将任务交给Dataextraction_agent处理，数据提取:包括记录查询(获取记录列表)数据导出(批量提取)、报表生成(格式化报表)。
                如果是其他问题，返回 other 。
                请根据用户输入的关键词进行分类，除了这几个选项外，请不要返回其他内容。请不要解释。
                """
    
    prompts = [
        {"role": "system", "content": prompt},
        {"role": "user", "content": state["messages"][0]}
    ]
    # 如果已经有type属性了，表示问题已经交由其他节点处理完成了，就可以直接返回
    if "type" in state:
        writer({"supervisor_step", f"已获得{state['type']}智能体处理结果"})
        return {"type": END}
    else:
        response = client.chat.completions.create(
            model="glm-4.5v",
            messages=prompts,
            temperature=0
        )
        typeRes = response.choices[0].message.content.strip()
        writer({"supervisor_step", f"问题分类结果：{typeRes}"})
        if typeRes in nodes:
            return {"type": typeRes}
        else:
            raise ValueError("type is not in (Supervisor,Imperativequery,ExploratoryAnalysis,Dataextraction,other)")
    return {}

def Imperativequery_node(state: State):
    print("Imperativequery node")
    writer = get_stream_writer()
    writer({"node": "Imperativequery_node"})

    return {"messages": [HumanMessage(content="Imperativequery_node")], "type": "Imperativequery"}

def ExploratoryAnalysis_node(state: State):
    print("ExploratoryAnalysis node")
    writer = get_stream_writer()
    writer({"node": "ExploratoryAnalysis_node"})

    return {"messages": [HumanMessage(content="ExploratoryAnalysis_node")], "type": "ExploratoryAnalysis"}

def Dataextraction_node(state: State):
    print("Dataextraction node")
    writer = get_stream_writer()
    writer({"node": "Dataextraction_node"})

    return {"messages": [HumanMessage(content="Dataextraction_node")], "type": "Dataextraction"}

#条件路由
def routing_func(state: State):
    if state["type"] == "Dataextraction":
        return "Dataextraction_node"
    elif state["type"] == "ExploratoryAnalysis":
        return "ExploratoryAnalysis_node"
    elif state["type"] == "Imperativequery":
        return "Imperativequery_node"
    elif state["type"] == END:
        return END
    else:
        return "other_node"


#构建图
builder = StateGraph(State)
#添加节点
builder.add_node("Supervisor_node", supervisor_node)
builder.add_node("Imperativequery_node", Imperativequery_node)
builder.add_node("ExploratoryAnalysis_node", ExploratoryAnalysis_node)#
builder.add_node("Dataextraction_node", Dataextraction_node)
builder.add_node("other_node", other_node)
#添加边（edge）
builder.add_edge(START, "Supervisor_node")
builder.add_conditional_edges("Supervisor_node", routing_func, ["Imperativequery_node","ExploratoryAnalysis_node","Dataextraction_node","other_node",END])
builder.add_edge("Imperativequery_node", "Supervisor_node")
builder.add_edge("ExploratoryAnalysis_node", "Supervisor_node")
builder.add_edge("Dataextraction_node", "Supervisor_node")
builder.add_edge("other_node", "Supervisor_node")

#创建流程图graph
checkpointer = InMemorySaver()
graph = builder.compile(checkpointer=checkpointer)

#执行任务测试代码
if __name__ == "__main__":
    config = {
        "configurable":{
            "thread_id": "1"
        }
    }

    for chunk in graph.stream({"messages":["你好"]}
                  ,config
                  ,stream_mode="custom"):
        print(chunk)

    # res = graph.invoke({"messages":["是哪些国家"]}
    #                    ,config
    #                    ,stream_mode="values")
    # print(res["messages"][-1].content)